clear
close all
clc

%params
siga=1;
sigv=2;
rhoa=.95;


%State variables
vaT=1;
va1t=2;
va0t=3;
a_t=4;
nend=4;

eaT=1;
vt=2;
nshocks=2;

GAM0=zeros(nend,nend);
GAM1=zeros(nend,nend);
PSI=zeros(nend,nshocks);

%\epsilon_{a,t+1}
GAM0(1,vaT)=1;
PSI(1,eaT)=siga;

%\epsilon^1_{a,t}
GAM0(2,va1t)=1;
GAM0(2,vaT)=-siga^2/(siga^2+sigv^2);
PSI(2,vt)=sigv;

%epsilon^0_t
GAM0(3,va0t)=1;
GAM1(3,vaT)=1;
GAM1(3,va1t)=-1;

%a_t
GAM0(4,a_t)=1;
GAM0(4,va0t)=-1;

GAM1(4,va1t)=1;
GAM1(4,a_t)=rhoa;



% Model solution
GG=pinv(GAM0)*GAM1;
RR=pinv(GAM0)*PSI;

nirf=10; ssirf=zeros(nend,nirf,2);
for j=1:2
    imp=zeros(2,1); imp(j)=1;
    for i=0:nirf
        ssirf(:,i+1,j)=GG^i*RR*imp;
    end
end
ssirf(abs(ssirf)<1e-6)=0;
figure
for i=1:4
    subplot(2,2,i)
    plot(0:nirf,ssirf(i,:,1))
    hold on
    plot(0:nirf,ssirf(i,:,1)*0)
    hold off
    axis tight
    if i==1
        title('Next period shock')
    elseif i==2
        title('Anticipated shock at t')
    elseif i==3
        title('Surprise shock at t')
    elseif i==4
        title('TFP process at t') 
    end
end
    figure
for i=1:4
    subplot(2,2,i)
    plot(0:nirf,ssirf(i,:,2))
    hold on
    plot(0:nirf,ssirf(i,:,2)*0)
    hold off
    axis tight
    if i==1
        title('Next period shock')
    elseif i==2
        title('Anticipated shock at t')
    elseif i==3
        title('Surprise shock at t')
    elseif i==4
        title('TFP process at t') 
    end
end